from io import StringIO

from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz
import pandas as pd
import pydotplus
from sklearn import model_selection
from sklearn import metrics


class PredictDiabetes(object):
    @classmethod
    def predict_diabetes(cls):
        data_set = pd.read_csv('pima-indians-diabetes.csv')
        # print(data_set)
        feature_columns = ['pregnant', 'glucose', 'bp', 'skin', 'insulin', 'bmi', 'pedigree', 'age']
        X = data_set[feature_columns]
        y = data_set.label

        dtc = DecisionTreeClassifier(criterion='entropy', max_depth=3)
        X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size=0.25)
        dtc.fit(X_train, y_train)
        y_pred = dtc.predict(X_test)
        y_score = metrics.accuracy_score(y_test, y_pred)
        print(y_score)

        dot_data = StringIO()
        export_graphviz(dtc, out_file=dot_data,
                        feature_names=feature_columns,
                        class_names=['0', '1'],
                        rounded=True,
                        filled=True,
                        special_characters=True)
        graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
        graph.write_pdf('decision_tree.pdf')


if __name__ == '__main__':
    PredictDiabetes.predict_diabetes()
